import sys
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                             QHBoxLayout, QPushButton, QLabel, QFileDialog,
                             QComboBox, QMessageBox, QSplitter)
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt
import cv2
import numpy as np


class ImageStitcher(QMainWindow):
    def __init__(self):
        super().__init__()
        self.image1 = None
        self.image2 = None
        self.image1_path = ""
        self.image2_path = ""
        self.stitched_result = None
        self.initUI()

    def initUI(self):
        self.setWindowTitle('图像拼接工具')
        self.setGeometry(100, 100, 1200, 800)

        # 创建主布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        # 两个图片上传区域
        upload_layout = QHBoxLayout()

        # 左侧上传区域
        left_upload_widget = QWidget()
        left_upload_layout = QVBoxLayout(left_upload_widget)

        self.select_button1 = QPushButton('选择第一张图像')
        self.select_button1.clicked.connect(lambda: self.select_image(1))

        self.image1_label = QLabel('第一张图像')
        self.image1_label.setAlignment(Qt.AlignCenter)

        self.image1_display = QLabel()
        self.image1_display.setAlignment(Qt.AlignCenter)
        self.image1_display.setMinimumSize(300, 300)
        self.image1_display.setStyleSheet("border: 1px solid #ccc;")

        left_upload_layout.addWidget(self.select_button1)
        left_upload_layout.addWidget(self.image1_label)
        left_upload_layout.addWidget(self.image1_display)

        # 右侧上传区域
        right_upload_widget = QWidget()
        right_upload_layout = QVBoxLayout(right_upload_widget)

        self.select_button2 = QPushButton('选择第二张图像')
        self.select_button2.clicked.connect(lambda: self.select_image(2))

        self.image2_label = QLabel('第二张图像')
        self.image2_label.setAlignment(Qt.AlignCenter)

        self.image2_display = QLabel()
        self.image2_display.setAlignment(Qt.AlignCenter)
        self.image2_display.setMinimumSize(300, 300)
        self.image2_display.setStyleSheet("border: 1px solid #ccc;")

        right_upload_layout.addWidget(self.select_button2)
        right_upload_layout.addWidget(self.image2_label)
        right_upload_layout.addWidget(self.image2_display)

        upload_layout.addWidget(left_upload_widget)
        upload_layout.addWidget(right_upload_widget)
        main_layout.addLayout(upload_layout)

        # 拼接设置区域
        settings_layout = QHBoxLayout()
        self.direction_label = QLabel('拼接方向:')
        self.direction_combo = QComboBox()
        self.direction_combo.addItems(['水平拼接', '垂直拼接'])

        self.clear_button = QPushButton('清除图像')
        self.clear_button.clicked.connect(self.clear_images)

        self.stitch_button = QPushButton('开始拼接')
        self.stitch_button.clicked.connect(self.stitch_images)
        self.stitch_button.setEnabled(False)

        settings_layout.addWidget(self.direction_label)
        settings_layout.addWidget(self.direction_combo)
        settings_layout.addStretch(1)
        settings_layout.addWidget(self.clear_button)
        settings_layout.addWidget(self.stitch_button)
        main_layout.addLayout(settings_layout)

        # 结果显示区域
        result_widget = QWidget()
        result_layout = QVBoxLayout(result_widget)

        self.result_label = QLabel('拼接结果')
        self.result_label.setAlignment(Qt.AlignCenter)

        self.result_display = QLabel()
        self.result_display.setAlignment(Qt.AlignCenter)
        self.result_display.setMinimumSize(600, 400)
        self.result_display.setStyleSheet("border: 1px solid #ccc;")

        self.save_button = QPushButton('保存结果')
        self.save_button.clicked.connect(self.save_result)
        self.save_button.setEnabled(False)

        result_layout.addWidget(self.result_label)
        result_layout.addWidget(self.result_display)
        result_layout.addWidget(self.save_button)
        main_layout.addWidget(result_widget)

        # 状态栏
        self.statusBar().showMessage('就绪')

    def select_image(self, image_index):
        file_path, _ = QFileDialog.getOpenFileName(
            self, f'选择第{image_index}张图像文件', '', '图像文件 (*.png *.jpg *.jpeg *.bmp)'
        )

        if file_path:
            try:
                img = cv2.imread(file_path)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                if image_index == 1:
                    self.image1 = img
                    self.image1_path = file_path
                    self.update_image_display(self.image1_display, img)
                else:
                    self.image2 = img
                    self.image2_path = file_path
                    self.update_image_display(self.image2_display, img)

                # 检查是否两张图片都已选择
                self.stitch_button.setEnabled(self.image1 is not None and self.image2 is not None)
                self.statusBar().showMessage(f'已加载第{image_index}张图像')
            except Exception as e:
                QMessageBox.critical(self, '错误', f'无法加载图像: {str(e)}')

    def clear_images(self):
        self.image1 = None
        self.image2 = None
        self.image1_path = ""
        self.image2_path = ""
        self.stitched_result = None

        self.image1_display.clear()
        self.image2_display.clear()
        self.result_display.clear()

        self.stitch_button.setEnabled(False)
        self.save_button.setEnabled(False)
        self.statusBar().showMessage('已清除所有图像')

    def update_image_display(self, display_widget, img):
        h, w, ch = img.shape
        bytes_per_line = ch * w
        q_img = QImage(img.data, w, h, bytes_per_line, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(q_img)
        display_widget.setPixmap(pixmap.scaled(
            display_widget.width(),
            display_widget.height(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        ))

    def stitch_images(self):
        if self.image1 is None or self.image2 is None:
            QMessageBox.warning(self, '警告', '请选择两张图像进行拼接')
            return

        try:
            direction = self.direction_combo.currentText()

            if direction == '水平拼接':
                # 寻找最大高度
                max_height = max(self.image1.shape[0], self.image2.shape[0])

                # 调整图像高度
                img1_padded = self.pad_image(self.image1, max_height, None)
                img2_padded = self.pad_image(self.image2, max_height, None)

                # 水平拼接
                self.stitched_result = np.hstack((img1_padded, img2_padded))
            else:  # 垂直拼接
                # 寻找最大宽度
                max_width = max(self.image1.shape[1], self.image2.shape[1])

                # 调整图像宽度
                img1_padded = self.pad_image(self.image1, None, max_width)
                img2_padded = self.pad_image(self.image2, None, max_width)

                # 垂直拼接
                self.stitched_result = np.vstack((img1_padded, img2_padded))

            self.update_image_display(self.result_display, self.stitched_result)
            self.save_button.setEnabled(True)
            self.statusBar().showMessage('拼接完成')

        except Exception as e:
            QMessageBox.critical(self, '错误', f'拼接失败: {str(e)}')
            self.statusBar().showMessage('拼接失败')

    def pad_image(self, img, target_height=None, target_width=None):
        h, w = img.shape[:2]

        if target_height is not None and h < target_height:
            padding = target_height - h
            return cv2.copyMakeBorder(img, 0, padding, 0, 0, cv2.BORDER_CONSTANT, value=[0, 0, 0])
        elif target_width is not None and w < target_width:
            padding = target_width - w
            return cv2.copyMakeBorder(img, 0, 0, 0, padding, cv2.BORDER_CONSTANT, value=[0, 0, 0])
        return img

    def save_result(self):
        if self.stitched_result is None:
            return

        file_path, _ = QFileDialog.getSaveFileName(
            self, '保存结果', '', '图像文件 (*.png *.jpg *.jpeg)'
        )

        if file_path:
            try:
                # 转换为BGR格式保存
                save_img = cv2.cvtColor(self.stitched_result, cv2.COLOR_RGB2BGR)
                cv2.imwrite(file_path, save_img)
                QMessageBox.information(self, '成功', f'图像已保存至: {file_path}')
                self.statusBar().showMessage(f'图像已保存至: {file_path}')
            except Exception as e:
                QMessageBox.critical(self, '错误', f'保存失败: {str(e)}')
                self.statusBar().showMessage('保存失败')


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = ImageStitcher()
    window.show()
    sys.exit(app.exec_())    